package com.lcj.ai.service;

import com.alibaba.cloud.ai.tongyi.audio.TongYiAudioSpeechOptions;
import com.alibaba.cloud.ai.tongyi.audio.api.SpeechPrompt;
import com.alibaba.cloud.ai.tongyi.audio.api.SpeechResponse;
import com.lcj.ai.model.ActorFilms;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.ChatClient;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.Generation;
import org.springframework.ai.chat.StreamingChatClient;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.image.ImageClient;
import org.springframework.ai.image.ImageOptions;
import org.springframework.ai.image.ImagePrompt;
import org.springframework.ai.image.ImageResponse;
import org.springframework.ai.model.ModelClient;
import org.springframework.ai.model.ModelOptions;
import org.springframework.ai.parser.BeanOutputParser;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import javax.annotation.processing.Completion;
import java.util.HashMap;
import java.util.Map;

/**
 * AI服务实现，提供多种AI功能接口，包括文本聊天、流式聊天、图片生成、语音生成等。
 */
@Slf4j
@Service
public class AiService {

    /**
     * 普通文本聊天客户端
     */
    private final ChatClient chatClient;
    /**
     * 流式文本聊天客户端
     */
    private final StreamingChatClient streamChatClient;
    /**
     * 图片生成客户端
     */
    private final ImageClient imageClient;
    /**
     * 语音生成客户端
     */
    private final ModelClient<SpeechPrompt, SpeechResponse> modelClient;

    /**
     * 构造函数，初始化各种AI客户端。
     */
    public AiService(ChatClient chatClient, StreamingChatClient streamChatClient, ImageClient imageClient,
                     ModelClient<SpeechPrompt, SpeechResponse> modelClient) {
        this.chatClient = chatClient;
        this.streamChatClient = streamChatClient;
        this.imageClient = imageClient;
        this.modelClient = modelClient;
    }

    /**
     * 普通聊天接口，发送消息并返回回复。
     *
     * @param message 用户发送的消息
     * @return 服务器回复的消息
     */
    public String normalCompletion(String message) {
        Prompt prompt = new Prompt(new UserMessage(message));
        return this.chatClient.call(prompt).getResult().getOutput().getContent();
    }

    /**
     * 流式聊天接口，发送消息并以流的形式返回回复。
     *
     * @param message 用户发送的消息
     * @return 响应流
     */
    public Flux<ChatResponse> streamCompletion(String message) {
        Prompt prompt = new Prompt(new UserMessage(message));
        return this.streamChatClient.stream(prompt);
    }

    /**
     * 生成图片接口，根据说明和规格生成图片。
     *
     * @param instructions 图片说明
     * @param options      图片生成选项
     * @return 图片生成响应
     */
    public ImageResponse genImage(String instructions, ImageOptions options) {
        ImagePrompt request = new ImagePrompt(instructions, options);
        return this.imageClient.call(request);
    }

    /**
     * 生成语音接口，根据文本和语音选项生成语音。
     *
     * @param instructions 语音内容
     * @param options      语音生成选项
     * @return 语音生成响应
     */
    public SpeechResponse genAudio(String instructions, ModelOptions options) {
        if (options instanceof TongYiAudioSpeechOptions speechOptions) {
            SpeechPrompt request = new SpeechPrompt(instructions, speechOptions);
            return this.modelClient.call(request);
        }
        return null;
    }

    /**
     * 获取演员的影片记录。
     *
     * @param actor 演员名称
     * @return 演员的影片记录
     */
    public ActorFilms getActorFilms(String actor) {
        var outputParser = new BeanOutputParser<>(ActorFilms.class);
        String format = outputParser.getFormat();
        log.info("格式: {}", format);
        String userMessage = """
                为演员{actor}生成影片记录.
                {format}
                """;
        PromptTemplate promptTemplate = new PromptTemplate(userMessage, Map.of("actor", actor, "format", format));
        Prompt prompt = promptTemplate.create();
        Generation generation = chatClient.call(prompt).getResult();
        // 解析生成的影片记录
        String content = generation.getOutput().getContent()
                .replace("```json", "")
                .replace("```", "");
        return outputParser.parse(content);
    }

    /**
     * 加载提示模板
     */
    @Value("classpath:/prompts/joke-prompt.st")
    private Resource jokeResource;

    /**
     * 加载并生成提示模板。
     *
     * @param adjective 形容词
     * @param topic     主题
     * @return 生成的提示模板消息
     */
    public AssistantMessage genPromptTemplates(String adjective, String topic) {
        PromptTemplate promptTemplate = new PromptTemplate(jokeResource);
        Prompt prompt = promptTemplate.create(Map.of("adjective", adjective, "topic", topic));
        return chatClient.call(prompt).getResult().getOutput();
    }

    /**
     * 外部资料
     */
    @Value("classpath:/docs/wikipedia-curling.md")
    private Resource docsToStuffResource;

    /**
     * 问答
     */
    @Value("classpath:/prompts/qa-prompt.st")
    private Resource qaPromptResource;

    /**
     * 问答接口，可选是否使用外部数据源。
     *
     * @param message   提问的内容
     * @param outerData 是否使用外部数据
     * @return 问题的回答
     */
    public Completion stuffCompletion(String message, boolean outerData) {
        PromptTemplate promptTemplate = new PromptTemplate(qaPromptResource);
        Map<String, Object> map = new HashMap<>();
        map.put("question", message);

        if (outerData) {
            map.put("context", docsToStuffResource);
        } else {
            map.put("context", "");
        }

        Prompt prompt = promptTemplate.create(map);
        Generation generation = chatClient.call(prompt).getResult();
        // 返回问题的回答
        return new Completion() {
            @Override
            public String getValue() {
                return generation.getOutput().getContent();
            }

            @Override
            public String getMessage() {
                return message;
            }
        };
    }
}
